package com.ag.improvedjava.iterable;

import java.util.Iterator;
import java.util.NoSuchElementException;

import com.ag.improvedjava.utils.Function;
import com.ag.improvedjava.utils.Function2;

final class SelectManyIterable<T, U, R> implements Iterable<R> {

	private Iterable<T> iterable;
	private Function<T, Iterable<U>> function;
	private Function2<T, U, R> resultSelector;
	
	SelectManyIterable(Iterable<T> iterable, Function<T, Iterable<U>> function) {
		this.iterable = iterable;
		this.function = function;
		this.resultSelector = new Function2<T, U, R>() {
			@SuppressWarnings("unchecked")
			@Override
			public R execute(T param1, U param2) {
				return (R)param2;
			}
		};
	}
	
	SelectManyIterable(Iterable<T> iterable, Function<T, Iterable<U>> function, Function2<T, U, R> resultSelector) {
		this.iterable = iterable;
		this.function = function;
		this.resultSelector = resultSelector; 
	}

	@Override
	public Iterator<R> iterator() {
		return new SelectManyIterator<>(iterable.iterator(), function, resultSelector);
	}
}

class SelectManyIterator<T, U, R> extends UnmodifiableIterator<R> {
	
	private final Iterator<T> iterator;
	private final Function<T, Iterable<U>> function;
	private final Function2<T, U, R> resultSelector;

	private T currentT;
	private Iterator<U> functionIterator;
		
	SelectManyIterator(Iterator<T> iterator, Function<T, Iterable<U>> function, Function2<T, U, R> resultSelector) {
		this.iterator = iterator;
		this.function = function;
		this.resultSelector = resultSelector;
	}

	@Override
	public boolean hasNext() {
		while (functionIterator == null || !functionIterator.hasNext()) {
			if (!iterator.hasNext()) {
				return false;
			}
			currentT = iterator.next();
			functionIterator = function.execute(currentT).iterator();
		}
		return true;
	}

	@Override
	public R next() {
		if (!hasNext()) {
			throw new NoSuchElementException();
		}
		return resultSelector.execute(currentT, functionIterator.next());		
	}

}
